package com.kuafuai.manage.mcp.server;

import com.fasterxml.jackson.core.type.TypeReference;
import com.kuafuai.manage.mcp.spec.McpError;
import com.kuafuai.manage.mcp.spec.McpSchema;
import com.kuafuai.manage.mcp.spec.McpServerSession;
import reactor.core.publisher.Mono;

import javax.servlet.AsyncContext;

/**
 * Represents an asynchronous exchange with a Model Context Protocol (MCP) client. The
 * exchange provides methods to interact with the client and query its capabilities.
 *
 * @author Dariusz Jędrzejczyk
 */
public class McpAsyncServerExchange {

    private final McpServerSession session;

    private final McpSchema.ClientCapabilities clientCapabilities;

    private final McpSchema.Implementation clientInfo;

    private static final TypeReference<McpSchema.CreateMessageResult> CREATE_MESSAGE_RESULT_TYPE_REF = new TypeReference<McpSchema.CreateMessageResult>() {
    };

    private static final TypeReference<McpSchema.ListRootsResult> LIST_ROOTS_RESULT_TYPE_REF = new TypeReference<McpSchema.ListRootsResult>() {
    };

    /**
     * Create a new asynchronous exchange with the client.
     * @param session The server session representing a 1-1 interaction.
     * @param clientCapabilities The client capabilities that define the supported
     * features and functionality.
     * @param clientInfo The client implementation information.
     */
    public McpAsyncServerExchange(McpServerSession session, McpSchema.ClientCapabilities clientCapabilities,
                                  McpSchema.Implementation clientInfo) {
        this.session = session;
        this.clientCapabilities = clientCapabilities;
        this.clientInfo = clientInfo;
    }

    /**
     * Get the client capabilities that define the supported features and functionality.
     * @return The client capabilities
     */
    public McpSchema.ClientCapabilities getClientCapabilities() {
        return this.clientCapabilities;
    }

    /**
     * Get the client implementation information.
     * @return The client implementation details
     */
    public McpSchema.Implementation getClientInfo() {
        return this.clientInfo;
    }

    /**
     * Create a new message using the sampling capabilities of the client. The Model
     * Context Protocol (MCP) provides a standardized way for servers to request LLM
     * sampling (“completions” or “generations”) from language models via clients. This
     * flow allows clients to maintain control over model access, selection, and
     * permissions while enabling servers to leverage AI capabilities—with no server API
     * keys necessary. Servers can request text or image-based interactions and optionally
     * include context from MCP servers in their prompts.
     * @param createMessageRequest The request to create a new message
     * @return A Mono that completes when the message has been created
     * @see McpSchema.CreateMessageRequest
     * @see McpSchema.CreateMessageResult
     * @see <a href=
     * "https://spec.modelcontextprotocol.io/specification/client/sampling/">Sampling
     * Specification</a>
     */
    public Mono<McpSchema.CreateMessageResult> createMessage(McpSchema.CreateMessageRequest createMessageRequest) {
        if (this.clientCapabilities == null) {
            return Mono.error(new McpError("Client must be initialized. Call the initialize method first!"));
        }
        if (this.clientCapabilities.getSampling() == null) {
            return Mono.error(new McpError("Client must be configured with sampling capabilities"));
        }
        return this.session.sendRequest(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, createMessageRequest,
                CREATE_MESSAGE_RESULT_TYPE_REF);
    }

    /**
     * Retrieves the list of all roots provided by the client.
     * @return A Mono that emits the list of roots result.
     */
    public Mono<McpSchema.ListRootsResult> listRoots() {
        return this.listRoots(null);
    }

    /**
     * Retrieves a paginated list of roots provided by the client.
     * @param cursor Optional pagination cursor from a previous list request
     * @return A Mono that emits the list of roots result containing
     */
    public Mono<McpSchema.ListRootsResult> listRoots(String cursor) {
        return this.session.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor),
                LIST_ROOTS_RESULT_TYPE_REF);
    }

    public AsyncContext getTransportAsyncContext(){
        return this.session.getTransport().getAsyncContext();
    }

}
